import torch
import torchvision
from torch import nn
from PIL import Image
import matplotlib.pyplot as plt
from torchvision.transforms import transforms


# 验证池化层参数stride函数
def pool_test_stride():
    max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
    input = torch.tensor([[1, 2, 3, 4, 5],
                          [6, 7, 8, 9, 10],
                          [11, 12, 13, 14, 15],
                          [16, 17, 18, 19, 20],
                          [21, 22, 23, 24, 25]])
    """
    PyTorch的卷积/池化层要求输入是一个4维张量，
    格式为：`(batch_size, channels, height, width)`，
    所以这里需要转化为4维张量。
    """
    input = torch.reshape(input, (-1, 1, 5, 5))
    print(f"输入input：{input}\n")
    output = max_pool(input)
    print(f"输出output：{output}\n")


# 验证ceil_mode参数
def pool_test_ceil_mode():
    max_pool_true = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
    max_pool_false = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=False)
    # 5*5的输入矩阵
    input = torch.tensor([[1, 2, 3, 4, 5],
                          [6, 7, 8, 9, 10],
                          [11, 12, 13, 14, 15],
                          [16, 17, 18, 19, 20],
                          [21, 22, 23, 24, 25]])
    input = torch.reshape(input, (-1, 1, 5, 5))
    print(f"输入input：{input}\n")
    output_true = max_pool_true(input)
    output_false = max_pool_false(input)
    print(f"输出output_true：{output_true}\n")
    print(f"输出output_false：{output_false}\n")


# 验证padding参数
def pool_test_padding():
    max_pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=1)
    input = torch.tensor([[1, 2, 3, 4, 5],
                          [6, 7, 8, 9, 10],
                          [11, 12, 13, 14, 15],
                          [16, 17, 18, 19, 20],
                          [21, 22, 23, 24, 25]])
    input = torch.reshape(input, (-1, 1, 5, 5))
    print(f"输入input：{input}\n")
    output = max_pool(input)
    print(f"输出output：{output}\n")


# 池化处理实际图像
def image_pool():
    """
    用最大池化处理实际图像，并可视化结果
    """
    # 1. 读取图像并转为灰度图（简化处理）
    img = Image.open("../images/dog.png").convert('L')  # 转为单通道灰度图
    plt.figure(figsize=(12, 6))

    # 2. 将图像转为Tensor并归一化
    transform = transforms.Compose([
        transforms.ToTensor(),  # 转为Tensor并自动归一化到[0,1]
    ])
    input_tensor = transform(img).unsqueeze(0)  # 增加batch维度 -> [1, 1, H, W]
    print(f"输入图像尺寸: {input_tensor.shape}")

    # 3. 定义池化层
    pool = nn.MaxPool2d(kernel_size=16, stride=16)  # 使用16x16的窗口，步长16

    # 4. 应用池化
    with torch.no_grad():
        output_tensor = pool(input_tensor)
    print(f"输出图像尺寸: {output_tensor.shape}")

    # 5. 转换为PIL图像并显示
    def tensor_to_image(tensor):
        return transforms.ToPILImage()(tensor.squeeze(0))  # 去掉batch维度

    # 原始图像
    plt.subplot(1, 2, 1)
    plt.imshow(img, cmap='gray')
    plt.title("Original Image")
    plt.axis('off')

    # 池化后的图像
    plt.subplot(1, 2, 2)  # Matplotlib 中用于创建子图（subplot）的函数，它的作用是在一个图像窗口中划分多个绘图区域，并指定当前操作的子图位置。
    output_img = tensor_to_image(output_tensor)
    plt.imshow(output_img, cmap='gray')
    plt.title(f"MaxPool (k=16, s=16)")
    plt.axis('off')

    plt.tight_layout()  # 自动调整子图之间的间距和布局
    plt.show()


if __name__ == '__main__':
    # pool_test_stride()
    # pool_test_ceil_mode()
    # pool_test_padding()
    image_pool()
